# Helper to align calendar‐time weights to event‐time
align_to_event <- function(w) {
  t0       <- which(is.na(w))[1]
  ev_index <- seq_along(w) - t0
  names(w) <- ev_index
  w[is.na(w)] <- 0
  w
}

# Compute combined event‐time weights for any cohorts
env_compute_combined_weights <- function(Lambda, T_Weights, ev_grid) {
  coh_cols     <- colnames(Lambda)
  cohort_years <- sub("^X", "", coh_cols)
  W_ev <- sapply(coh_cols, function(col) {
    w_cal <- Lambda[, col]
    w_ev  <- align_to_event(w_cal)
    w_sub <- w_ev[ev_grid]
    w_sub[is.na(w_sub)] <- 0
    w_sub
  })
  colnames(W_ev) <- cohort_years
  weights_vec <- T_Weights[cohort_years]
  combined    <- as.numeric(W_ev %*% weights_vec)
  names(combined) <- ev_grid
  shifted <- c(0, combined[-length(combined)])
  names(shifted) <- ev_grid
  shifted
}

# Single‐panel event‐study + time‐weights plot with optional FEct and dynamic legend
plot_event_with_weights <- function(
    sdid_x, sdid_est, sdid_se,
    Lambda, T_Weights,
    main,
    fect_x = NULL, fect_est = NULL, fect_se = NULL,
    fect_offset = -0.12, sdid_offset = +0.12,
    xlab = "Elections Relative to First Election With Strict Voter ID Law",
    ylab = "Difference in Voter Turnout Percentage (%)",
    xlim = range(sdid_x), ylim = c(-10, 10)
) {
  # Colors
  c_fect <- "#0A8754"
  c_sdid <- "#AF3B6E"
  
  x_seq  <- sort(unique(sdid_x))
  ev_grid <- as.character(x_seq)
  
  # compute shading weights
  combined_w <- env_compute_combined_weights(Lambda, T_Weights, ev_grid)
  
  # blank canvas
  plot(xlim, ylim, type = "n", xlab = xlab, ylab = ylab,
       xaxt = "n", yaxt = "n")
  abline(h = 0,   col = "gray50", lwd = 1.5, lty = 2)
  abline(v = 0.5, col = "gray50", lwd = 1.5, lty = 2)
  title(main = main, line = 0.5, cex.main = 1.2)
  
  usr <- par("usr"); base_y <- usr[3]; h <- (usr[4] - usr[3]) * 0.10
  x0      <- x_seq[x_seq <= 0]
  y_band  <- base_y + (combined_w[as.character(x0)] / max(combined_w)) * h
  x_shade <- c(x0, 0.5)
  y_shade <- c(y_band, tail(y_band, 1))
  polygon(
    c(x_shade, rev(x_shade)),
    c(y_shade, rep(base_y, length(x_shade))),
    col    = "#C492B1",
    border = "black",
    lwd    = 1
  )
  
  # Determine if FEct is present
  has_fect <- !is.null(fect_est) && length(fect_est) > 0
  
  # SDID points + CIs
  sdid_col <- if (has_fect) c_sdid else "black"
  points(sdid_x + sdid_offset, sdid_est, pch = 15, col = sdid_col, cex = 1)
  arrows(sdid_x + sdid_offset, sdid_est - 1.96 * sdid_se,
         sdid_x + sdid_offset, sdid_est + 1.96 * sdid_se,
         angle = 90, code = 3, length = 0.04, col = sdid_col)
  
  # Optional FEct points + CIs
  if (has_fect) {
    points(fect_x + fect_offset, fect_est, pch = 16, col = c_fect, cex = 1.2)
    arrows(fect_x + fect_offset, fect_est - 1.96 * fect_se,
           fect_x + fect_offset, fect_est + 1.96 * fect_se,
           angle = 90, code = 3, length = 0.03, col = c_fect)
  }
  
  # Axes
  axis(1, at = x_seq,    labels = x_seq)
  axis(2, at = seq(ylim[1], ylim[2], by = 2), labels = seq(ylim[1], ylim[2], by = 2))
  
  # Legend if FEct present
  if (has_fect) {
    legend("topleft",
           legend = c("FEct", "SDID"),
           col    = c(c_fect, c_sdid),
           pch    = c(16, 15),
           bty    = "n",
           cex    = 1.25)
  }
}
# ─── generate figures ───────────────────────────────────────────────────────
# comparison.eps
postscript("figures/comparison.eps", width = 12, height = 10,
           horizontal = FALSE, paper = "special")
layout(matrix(c(1,1,2,3), nrow = 2, byrow = TRUE), widths = c(1,1), heights = c(1,1))
par(mgp = c(1.5, 0.5, 0), mar = c(2.5,2.5,1.5,0))
# All
plot_event_with_weights(
  fect_x   = as.numeric(names(all_fect$est.att[,1])), fect_est = all_fect$est.att[,1], fect_se = all_fect$est.att[,2],
  sdid_x   = -15:8,                                   sdid_est = all_sdid_es$Estimate,       sdid_se = all_sdid_es$SE,
  Lambda   = all_sdid_es$Lambda, T_Weights = all_sdid_es$T_Weights,
  main     = "All Election Types Combined",
  fect_offset = -0.14, sdid_offset = +0.14,
  xlim     = c(-15,8),
  ylim = c(-12, 12)
)
# Presidential
plot_event_with_weights(
  fect_x   = as.numeric(names(presidential_fect$est.att[,1])), fect_est = presidential_fect$est.att[,1], fect_se = presidential_fect$est.att[,2],
  sdid_x   = -7:4,                                           sdid_est = presidential_sdid_es$Estimate, sdid_se = presidential_sdid_es$SE,
  Lambda   = presidential_sdid_es$Lambda, T_Weights = presidential_sdid_es$T_Weights,
  main     = "Presidential Elections",
  fect_offset = -0.12, sdid_offset = +0.12,
  xlim     = c(-7,4),
  ylim = c(-12, 12)
)
# Midterm
plot_event_with_weights(
  fect_x   = as.numeric(names(midterm_fect$est.att[,1])), fect_est = midterm_fect$est.att[,1], fect_se = midterm_fect$est.att[,2],
  sdid_x   = -7:4,                                        sdid_est = midterm_sdid_es$Estimate,     sdid_se = midterm_sdid_es$SE,
  Lambda   = midterm_sdid_es$Lambda, T_Weights = midterm_sdid_es$T_Weights,
  main     = "Midterm Elections",
  fect_offset = -0.12, sdid_offset = +0.12,
  ylab     = "",
  xlim     = c(-7,4),
  ylim = c(-12, 12)
)
dev.off()

# early_late.eps (only SDID panels with weights)
postscript("figures/early_late.eps", width = 14, height = 9,
           horizontal = FALSE, paper = "special")
par(oma = c(0,1.75,1.75,0))
layout(matrix(c(1,1,3,5, 2,2,4,6), nrow = 2, byrow = TRUE), widths = c(1,1), heights = c(1,1))

par(mgp = c(1.5,0.5,0), mar = c(2.5,2.6,0,1.5))
# All Early SDID
plot_event_with_weights(
  fect_x   = numeric(0), fect_est = numeric(0), fect_se = numeric(0),  # no FEct
  sdid_x   = -11:8,       sdid_est = all_early_sdid_es$Estimate, sdid_se = all_early_sdid_es$SE,
  Lambda   = all_early_sdid_es$Lambda, T_Weights = all_early_sdid_es$T_Weights,
  main     = "All Election Types",
  fect_offset = 0, sdid_offset = 0,
  xlab     = "",
  ylab     = "Difference in Voter Turnout Percentage (%)",
  xlim     = c(-11,8)
)
# Col 1
# Add column header for first column
mtext("All Election Types", side = 3, line = .4, outer = TRUE, at = 0.25, cex = 1, font = 2)
# All Late SDID
par(mgp = c(1.5,0.5,0), mar = c(2.5,2.6,0,1.5))
plot_event_with_weights(
  fect_x   = numeric(0), fect_est = numeric(0), fect_se = numeric(0),
  sdid_x   = -15:5,       sdid_est = all_late_sdid_es$Estimate,  sdid_se = all_late_sdid_es$SE,
  Lambda   = all_late_sdid_es$Lambda, T_Weights = all_late_sdid_es$T_Weights,
  main     = "",
  xlab     = "Elections Relative to First Election With Strict Voter ID Law",
  ylab     = "Difference in Voter Turnout Percentage (%)",
  xlim     = c(-15,5)
)
# Col 2
# Add row label for first plot of row 2
mtext("Early Adopters (before 2008)", side = 2, line = 0.4, outer = TRUE, at = 0.75, cex = 1, font = 2)
# Add column header for second column
mtext("Presidential Elections", side = 3, line = 0.4, outer = TRUE, at = 0.625, cex = 1, font = 2)
names(presidential_early_sdid_es$T_Weights) = "2008"
names(presidential_early_sdid_es$Lambda) = "2008"
# Presidential Early
par(mgp = c(1.5,0.5,0), mar = c(2.5,1.5,0,0.75))
plot_event_with_weights(
  fect_x   = numeric(0),
  sdid_x   = -5:4,      sdid_est = presidential_early_sdid_es$Estimate, sdid_se = presidential_early_sdid_es$SE,
  Lambda   = presidential_early_sdid_es$Lambda, T_Weights = presidential_early_sdid_es$T_Weights,
  main     = "Presidential Elections",
  xlab     = "",
  ylab     = "",
  xlim     = c(-5,4)
)
# Presidential Late
par(mgp = c(1.5,0.5,0), mar = c(2.5,1.5,0,0.75))
plot_event_with_weights(
  fect_x   = numeric(0),
  sdid_x   = -7:3,      sdid_est = presidential_late_sdid_es$Estimate, sdid_se = presidential_late_sdid_es$SE,
  Lambda   = presidential_late_sdid_es$Lambda, T_Weights = presidential_late_sdid_es$T_Weights,
  main     = "",
  xlab     = "Elections Relative to First Election With Strict Voter ID Law",
  ylab     = "",
  xlim     = c(-7,3)
)
# Col 3
# Add row label for first plot of row 3
mtext("Late Adopters (after 2008)", side = 2, line = 0.4, outer = TRUE, at = 0.25, cex = 1, font = 2)
# Add column header for third column
mtext("Midterm Elections", side = 3, line = 0.4, outer = TRUE, at = 0.875, cex = 1, font = 2)

# Midterm Early
par(mgp = c(1.5,0.5,0), mar = c(2.5,2.25,0,0.1))
plot_event_with_weights(
  fect_x   = numeric(0),
  sdid_x   = -5:4,      sdid_est = midterm_early_sdid_es$Estimate, sdid_se = midterm_early_sdid_es$SE,
  Lambda   = midterm_early_sdid_es$Lambda, T_Weights = midterm_early_sdid_es$T_Weights,
  main     = "Midterm Elections",
  xlab     = "",
  ylab     = "",
  xlim     = c(-5,4)
)
# Midterm Late
par(mgp = c(1.5,0.5,0), mar = c(2.5,2.25,0,0.1))
plot_event_with_weights(
  fect_x   = numeric(0),
  sdid_x   = -7:2,      sdid_est = midterm_late_sdid_es$Estimate, sdid_se = midterm_late_sdid_es$SE,
  Lambda   = midterm_late_sdid_es$Lambda, T_Weights = midterm_late_sdid_es$T_Weights,
  main     = "",
  xlab     = "Elections Relative to First Election With Strict Voter ID Law",
  ylab     = "",
  xlim     = c(-7,2)
)

dev.off()


